import pybullet as p

class ManualController(object):
    def __init__(self, controller='keyboard', outputRange=[[-200, 200], [-1.0, 1.0]], sensitivity=0.3, dt=1/240):
        self.type = controller
        self.outputRange = outputRange
        self.sensitivity = sensitivity
        self.cmd_raw = [0.75, 0.5]
        self.dt = dt
        if controller == 'xbox-linux':
            from xbox360controller import Xbox360Controller
            self.xbox_linux = Xbox360Controller(0, axis_threshold=0.2)
        if controller == 'xbox-general':
            import xbox360_controller
            self.xbox_general = xbox360_controller.Controller(0)

    def updateCommand(self):
        if self.type == 'keyboard':
            keyUp = p.B3G_UP_ARROW
            keyDown = p.B3G_DOWN_ARROW
            keyLeft = p.B3G_LEFT_ARROW
            keyRight = p.B3G_RIGHT_ARROW

            keys = p.getKeyboardEvents()
            if keyUp in keys and keys[keyUp]&p.KEY_IS_DOWN:
                self.cmd_raw[0] += self.sensitivity * self.dt
            if keyDown in keys and keys[keyDown]&p.KEY_IS_DOWN:
                self.cmd_raw[0] -= self.sensitivity * self.dt
            if keyRight in keys and keys[keyRight]&p.KEY_IS_DOWN:
                self.cmd_raw[1] += self.sensitivity * self.dt
            if keyLeft in keys and keys[keyLeft]&p.KEY_IS_DOWN:
                self.cmd_raw[1] -= self.sensitivity * self.dt
            
            if keyUp not in keys and keyDown not in keys:
                self.cmd_raw[0] = self.cmd_raw[0] * 0.97 + 0.75 * 0.03
            if keyRight not in keys and keyLeft not in keys:
                self.cmd_raw[1] = self.cmd_raw[1] * 0.97 + 0.5 * 0.03
            
            self.cmd_raw = [max(min(x, 1.0), 0.0) for x in self.cmd_raw]
        if self.type == 'xbox-linux':
            print(f'axis l: {self.xbox_linux.axis_l.y}; axis r: {self.xbox_linux.axis_r.x}')
            self.cmd_raw = [- self.xbox_linux.axis_l.y * 0.5 + 0.5, self.xbox_linux.axis_r.x * 0.5 + 0.5]
        if self.type == 'xbox-general':
            left_x, left_y = self.xbox_general.get_left_stick()
            right_x, right_y = self.xbox_general.get_right_stick()
            self.cmd_raw = [- left_y * 0.5 + 0.5, right_x * 0.5 + 0.5]
        # print(f'up cmd raw: {self.cmd_raw[0]}; rot cmd raw: {self.cmd_raw[1]}')
        force = self.outputRange[0][0] * (1.0 - self.cmd_raw[0]) + self.outputRange[0][1] * self.cmd_raw[0]
        torque = self.outputRange[1][0] * (1.0 - self.cmd_raw[1]) + self.outputRange[1][1] * self.cmd_raw[1]
        print(f'upward force: {force}; torque: {torque}')
        return [force / 2 - torque / 0.3, force / 2 + torque / 0.3]